import sys
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QGridLayout

class ArmController(QWidget):
    
    def __init__(self):
        super().__init__()
        self.controller = None
        self.jointAngles1 = [0.0] * 6
        self.jointAngles2 = [0.0] * 4
        self.speed = 50.0
        self.initUI()

    def initUI(self):
        self.setWindowTitle('Arm Controller')
        self.setGeometry(100, 100, 600, 400)

        layout = QVBoxLayout()

        connectButton = QPushButton('Connect Controller', self)
        connectButton.clicked.connect(self.connectController)
        layout.addWidget(connectButton)

        armsLayout = QHBoxLayout()

        arm1Layout = QVBoxLayout()
        arm1Layout.addWidget(QLabel('<b>Arm 1 (6 Joints)</b>', alignment=Qt.AlignCenter))

        grid1 = QGridLayout()
        self.angleLineEdits1 = []
        for i in range(6):
            label = QLabel(f'Joint {i+1}:', self)
            lineEdit = QLineEdit(str(self.jointAngles1[i]), self)
            lineEdit.setAlignment(Qt.AlignRight)
            lineEdit.textChanged.connect(lambda text, idx=i: self.updateJointAngle(text, idx, 1))
            grid1.addWidget(label, i, 0)
            grid1.addWidget(lineEdit, i, 1)
            self.angleLineEdits1.append(lineEdit)

        incrementButton1 = QPushButton('Increment by 5', self)
        incrementButton1.clicked.connect(lambda: self.incrementJoints(1))
        moveButton1 = QPushButton('Move Arm 1', self)
        moveButton1.clicked.connect(lambda: self.moveArm(1))
        buttonLayout1 = QHBoxLayout()
        buttonLayout1.addWidget(incrementButton1)
        buttonLayout1.addWidget(moveButton1)
        arm1Layout.addLayout(grid1)
        arm1Layout.addLayout(buttonLayout1)

        arm2Layout = QVBoxLayout()
        arm2Layout.addWidget(QLabel('<b>Arm 2 (4 Joints)</b>', alignment=Qt.AlignCenter))

        grid2 = QGridLayout()
        self.angleLineEdits2 = []
        for i in range(4):
            label = QLabel(f'Joint {i+7}:', self)
            lineEdit = QLineEdit(str(self.jointAngles2[i]), self)
            lineEdit.setAlignment(Qt.AlignRight)
            lineEdit.textChanged.connect(lambda text, idx=i: self.updateJointAngle(text, idx, 2))
            grid2.addWidget(label, i, 0)
            grid2.addWidget(lineEdit, i, 1)
            self.angleLineEdits2.append(lineEdit)

        incrementButton2 = QPushButton('Increment by 5', self)
        incrementButton2.clicked.connect(lambda: self.incrementJoints(2))
        moveButton2 = QPushButton('Move Arm 2', self)
        moveButton2.clicked.connect(lambda: self.moveArm(2))
        buttonLayout2 = QHBoxLayout()
        buttonLayout2.addWidget(incrementButton2)
        buttonLayout2.addWidget(moveButton2)
        arm2Layout.addLayout(grid2)
        arm2Layout.addLayout(buttonLayout2)

        armsLayout.addLayout(arm1Layout)
        armsLayout.addLayout(arm2Layout)

        speedLayout = QHBoxLayout()
        speedLabel = QLabel('Speed (mm/s):', self)
        self.speedLineEdit = QLineEdit(str(self.speed), self)
        self.speedLineEdit.setAlignment(Qt.AlignRight)
        self.speedLineEdit.textChanged.connect(self.updateSpeed)
        speedLayout.addWidget(speedLabel)
        speedLayout.addWidget(self.speedLineEdit)

        layout.addLayout(armsLayout)
        layout.addLayout(speedLayout)

        self.setLayout(layout)

    def connectController(self):
        if not self.controller:
            from driver_server import VSRController  # 假设VSRController可以从main.py导入
            self.controller = VSRController(simulate=True)
            self.controller.connect()
            print("Controller connected.")

    def closeEvent(self, event):
        """重写关闭事件处理"""
        if self.controller:
            self.controller.stop()
            print("Controller disconnected.")
        event.accept()

    def updateJointAngle(self, text, index, armIndex):
        try:
            angle = float(text)
            if armIndex == 1:
                self.jointAngles1[index] = angle
            elif armIndex == 2:
                self.jointAngles2[index] = angle
        except ValueError:
            pass

    def incrementJoints(self, armIndex):
        if armIndex == 1:
            self.jointAngles1 = [angle + 5 for angle in self.jointAngles1]
        elif armIndex == 2:
            self.jointAngles2 = [angle + 5 for angle in self.jointAngles2]
        self.updateLineEdits()

    def updateLineEdits(self):
        for i, lineEdit in enumerate(self.angleLineEdits1):
            lineEdit.setText(str(self.jointAngles1[i]))
        for i, lineEdit in enumerate(self.angleLineEdits2):
            lineEdit.setText(str(self.jointAngles2[i]))

    def updateSpeed(self, text):
        try:
            self.speed = float(text)
        except ValueError:
            pass

    def moveArm(self, armIndex):
        if not self.controller:
            print("Please connect the controller first.")
            return

        targetPose = self.jointAngles1 + self.jointAngles2 if armIndex == 1 else self.jointAngles1 + self.jointAngles2
        self.controller.set_joint_angle(targetPose, self.speed)



if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = ArmController()
    ex.show()
    sys.exit(app.exec_())


